# encoding=utf8

from __future__ import unicode_literals
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.feature_extraction.text import TfidfTransformer
import pickle
import sys

if sys.version < '3':
    reload(sys)
    sys.setdefaultencoding('utf8')


if __name__ == '__main__':
    data_file = sys.argv[1]
    v_file = sys.argv[2]
    tfidf_file = sys.argv[3]
    cat_count_file = sys.argv[4]
    cat_count = int(sys.argv[5])
    out_file = sys.argv[6]

    target = []
    X = []

    cat_dict = dict()
    with open(cat_count_file) as fd:
        for l in fd:
            d = l.strip().split(' ')
            if len(d) == 2:
                try:
                    count, cat = int(d[0]), d[1]
                    cat_dict[cat] = count

                except Exception as e:
                    print(str(e))

    print('cat_count is : ', len([k for k in cat_dict.keys() if cat_dict[k] >= cat_count]))
    with open(data_file) as fd:
        for l in fd:
            d = l.strip().split('\t')
            if len(d) == 2:
                t, x = d
                if cat_dict.get(t) and cat_dict[t] >= cat_count and t != '#18':
                    target.append(t.strip())
                    X.append(' '.join([k.strip() for k in x.strip().split(' ') if (not(k.strip().isdigit())) and 0 < len(k.strip()) < 15]))

    CV = CountVectorizer(min_df=5, ngram_range=(1, 1))
    X = CV.fit_transform(X)
    # print(X[0:10])
    TFIDF = TfidfTransformer()
    X = TFIDF.fit_transform(X)

    with open(v_file, 'wb+') as fd:
        pickle.dump(CV, fd)

    with open(tfidf_file, 'wb+') as fd:
        pickle.dump(TFIDF, fd)

    with open(out_file, 'wb+') as fd:
        pickle.dump((X, target), fd)
